#include <sys/mman.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <netinet/in.h>
#include <arpa/inet.h>

#include <endian.h>
#include <errno.h>
#include <fcntl.h>
#include <netdb.h>
#include <pthread.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <list>
#include <string>
#include <sstream>

#define exiterr(format, ...)\
do\
{\
    fprintf(stderr, format, ## __VA_ARGS__ );\
    exit(EXIT_FAILURE);\
} while(0)

#define liberr(name)\
do\
{\
    perror(#name);\
    exit(EXIT_FAILURE);\
} while(0)

#define safe_call(name, ...)\
do\
{\
    int err = name ( __VA_ARGS__ );\
    if(err)\
        liberr(name);\
} while(0)

std::list<uint16_t> available_ports;

pthread_mutex_t port_lock;

std::string listen_addr;

__thread int* counter;

void init_ports(uint16_t port_min, uint16_t port_max)
{
    for(uint16_t port = port_min; port < port_max; ++port)
        available_ports.push_back(port);
    safe_call(pthread_mutex_init, &port_lock, NULL);
}

bool get_ports(size_t n_ports, uint16_t ports[])
{
    safe_call(pthread_mutex_lock, &port_lock);
    
    bool result;
    if(available_ports.size() < n_ports)
        result = false;
    else
    {
        for(size_t i = 0; i < n_ports; ++i)
        {
            ports[i] = available_ports.front();
            available_ports.pop_front();
        }
        result = true;
    }
    
    safe_call(pthread_mutex_unlock, &port_lock);
    
    return result;
}

void return_ports(size_t n_ports, uint16_t const ports[])
{
    safe_call(pthread_mutex_lock, &port_lock);
    
    for(size_t i = 0; i < n_ports; ++i)
        available_ports.push_back(ports[i]);
    
    safe_call(pthread_mutex_unlock, &port_lock);
}

char const * const run_path = "client/client";

void ignored(ssize_t){}

// protocol:
//      in  byte 1: clients needed
//      out byte n * 2: ports, network byte order
//      
//      Port 0 is sent on error conditions
void* serve(void *connfd_v)
{
    int connfd = (int)(long)connfd_v;
    
    fputs("Connection established.\n", stderr);
    
    uint8_t n_ports;
    
    uint16_t zero = 0;
    
    // Initialize the per-thread counter
    int zerofd = open("/dev/zero", O_RDWR);
    if(zerofd < 0)
        liberr(open);
    counter = static_cast<int*>(mmap(NULL, sizeof(*counter), PROT_READ | PROT_WRITE, MAP_SHARED | MAP_ANONYMOUS, zerofd, 0));
    if(counter == NULL)
        liberr(mmap);
    safe_call(close, zerofd);
    *counter = 0;
    
    if(read(connfd, &n_ports, 1) == 1)
    {
        pid_t pid[n_ports];
        uint16_t ports[n_ports];
        
        memset(pid, 0, sizeof(pid_t) * n_ports);
        memset(ports, 0, 2 * n_ports);
        
        if(n_ports == 0 || !get_ports(n_ports, ports))
            ignored(write(connfd, &zero, 2));
        else
        {
            for(size_t i = 0; i < n_ports; ++i)
            {
                pid[i] = fork();
                if(pid[i] < 0)
                {
                    if(i > 0)
                        kill(-pid[0], SIGTERM);
                    ignored(write(connfd, &zero, 2));
                    goto clean;
                }
                else if(pid[i] == 0)
                {
                    if(i == 0)
                        setpgid(0, 0);
                    else
                        setpgid(0, pid[0]);
                    __sync_add_and_fetch(counter, 1);
                    fprintf(stderr, "In group: %d\n", (int)getpgid(0));
                    char listen_argument[32];
                    snprintf(listen_argument, 32, "%s:%u", listen_addr.c_str(), (unsigned)ports[i]);
                    execlp(run_path, run_path, "--eq-client", "--eq-listen", listen_argument, NULL);
                    liberr(execlp);
                }
            }
            
            while(true)
            {
                __sync_synchronize();
                if(*counter == n_ports)
                    break;
            }
            
            uint16_t ports_send[n_ports];
            for(size_t i = 0; i < n_ports; ++i)
                ports_send[i] = htons(ports[i]);
            
            if(write(connfd, ports_send, 2 * n_ports) == 2 * n_ports)
            {
                static uint8_t trashcan[4096];
                
                while(read(connfd, trashcan, 4096) > 0)
                    ;
                perror("read");
            }
            
            fprintf(stderr, "Killing group: %d\n", (int)pid[0]);
            kill(-pid[0], SIGTERM);
            clean:
            return_ports(n_ports, ports);
        }
    }
    
    safe_call(close, connfd);
    safe_call(munmap, counter, sizeof(*counter));
    
    fputs("Connection terminated.\n", stderr);
    
    return NULL;
}

void chld_handler(int)
{
    while(true)
    {
        pid_t dead_child = waitpid(-1, NULL, WNOHANG);
        if( dead_child == 0 || (dead_child < 0 && errno == ECHILD) )
            return;
        else if(dead_child < 0)
            liberr(waitpid);
    }
}

void setup_signal()
{
    struct sigaction chld_action;
    
    chld_action.sa_handler = chld_handler;
    safe_call(sigemptyset, &chld_action.sa_mask);
    chld_action.sa_flags = SA_NOCLDSTOP | SA_NODEFER | SA_RESTART;
    
    safe_call(sigaction, SIGCHLD, &chld_action, NULL);
}

char const* usage = "Usage: %s ip:port port_begin port_end\n";

int main(int argc, char *argv[])
{
    if(argc != 4)
        exiterr(usage, argv[0]);
    
    uint16_t port_begin, port_end;
    
    std::istringstream argument(argv[1]);
    std::string hostname, service;
    if(!std::getline(argument, hostname, ':') || !std::getline(argument, service, ':'))
        exiterr(usage, argv[0]);
    
    argument.clear();
    argument.str(argv[2]);
    if(!(argument >> port_begin))
        exiterr(usage, argv[0]);
    
    argument.clear();
    argument.str(argv[3]);
    if(!(argument >> port_end))
        exiterr(usage, argv[0]);
    
    addrinfo *listen_addrs = NULL;
    
    int err = getaddrinfo(hostname.c_str(), service.c_str(), NULL, &listen_addrs);
    if(err)
        exiterr("getaddrinfo: %s\n", gai_strerror(err));
    
    int listenfd = socket(listen_addrs->ai_family, SOCK_STREAM | SOCK_CLOEXEC, 0);
    if(listenfd < 0)
        liberr(socket);
    
    safe_call(bind, listenfd, listen_addrs->ai_addr, listen_addrs->ai_addrlen);
    
    freeaddrinfo(listen_addrs);
    
    safe_call(listen, listenfd, 5);
    
    if(listen_addrs->ai_addr->sa_family != AF_INET)
    {
        fputs("Error: listen address is not an IPv4 address.\n", stderr);
        return EXIT_FAILURE;
    }
    listen_addr = inet_ntoa(reinterpret_cast<sockaddr_in*>(listen_addrs->ai_addr)->sin_addr);
    
    setup_signal();
    init_ports(port_begin, port_end);
    
    int connfd;
    while(connfd = accept4(listenfd, NULL, NULL, SOCK_CLOEXEC), connfd >= 0)
    {
        pthread_t tid;
        safe_call(pthread_create, &tid, NULL, serve, (void*)connfd);
        safe_call(pthread_detach, tid);
    }
    
    liberr(accept);
}
